package com.opencv;

import org.opencv.core.*;
import org.opencv.imgcodecs.Imgcodecs;
import org.opencv.imgproc.Imgproc;
import pers.wz.framework.common.util.CommonUtil;
import pers.wz.framework.common.util.LogConsole;
import pers.wz.framework.common.util.ThrowableUtil;
import pers.wz.framework.common.util.awt.AwtUtil;
import pers.wz.framework.common.util.io.IoUtil;
import pers.wz.framework.common.util.lang.ExpectedException;

import java.awt.image.BufferedImage;
import java.io.File;
import java.io.FileOutputStream;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import static pers.wz.framework.common.util.lang.SystemProperties.projectAbsolutePath;

public class OpencvApi {



    public static void main(String[] args) {
        init(2, null);
//        System.out.println(compareHist("D:\\develop\\JavaCode\\pers-wz-soft\\opencv\\5d723d62-51e8-42b6-bd5e-a854b19bd5c0.png",
//                "D:\\develop\\JavaCode\\pers-wz-soft\\opencv\\646dfb1c-b583-47b5-aafa-325355d1b3b1.png", Imgproc.CV_COMP_CORREL));

        MatchResult matchResult = matchTemplate("images\\win10.png", "images\\Chrome.png", Imgproc.TM_CCORR_NORMED);




        System.out.println(matchResult.getSimilarity());
    }









    private static final AtomicBoolean INITIALIZED = new AtomicBoolean();
    private static String oprDir;

    /**
     * 初始化类库
     * @param osBit 操作系统位数 1-32位 2-64位
     * @param oprDir opencv工作的临时存储目录
     * */
    public static void init(int osBit, String oprDir) {
        try {
            if (osBit == 1) {
                System.load(projectAbsolutePath + "/dll/opencv/opencv_java451_x86.dll");
            } else {
                System.load(projectAbsolutePath + "/dll/opencv/opencv_java451_x64.dll");
            }
            if (oprDir != null) {
                OpencvApi.oprDir = oprDir;
            } else {
                OpencvApi.oprDir = projectAbsolutePath + "/opencv";
            }
            File file = new File(OpencvApi.oprDir);
            if (!file.exists()) {
                file.mkdirs();
            }
            INITIALIZED.set(true);
        } catch (Throwable e) {
            LogConsole.logError(OpencvApi.class, "init dll error: " + ThrowableUtil.getFullInfo(e));
            INITIALIZED.set(false);
        }
    }

    /**
     * 检查opencv类库是否正确加载
     * */
    private static void checkInitialized() {
        if (!INITIALIZED.get()) {
            throw new ExpectedException("OpencvApi is not initialized");
        }
    }

    /**
     * 获取两张图片相似度
     * @param image1 第一张图片
     * @param image2 第二张图片
     * @param method 比较方法
     *               Imgproc.CV_COMP_CORREL, CV_COMP_CHISQR, CV_COMP_INTERSECT, CV_COMP_BHATTACHARYYA
     *               一般选择CV_COMP_CORREL
     * */
    public static double compareHist(BufferedImage image1, BufferedImage image2, int method) {
        checkInitialized();
        String imagePath1 = oprDir + File.separator + UUID.randomUUID().toString() + ".png";
        String imagePath2 = oprDir + File.separator + UUID.randomUUID().toString() + ".png";
        try (OutputStream outputStream1 = new FileOutputStream(imagePath1);
             OutputStream outputStream2 = new FileOutputStream(imagePath2)) {
            AwtUtil.writeBufferedImage(image1, "png", outputStream1);
            AwtUtil.writeBufferedImage(image2, "png", outputStream2);
        } catch (Throwable e) {
            IoUtil.delete(new File(imagePath1));
            IoUtil.delete(new File(imagePath2));
            LogConsole.logError(OpencvApi.class, ThrowableUtil.getFullInfo(e));
            return 0;
        }
        try {
            return compareHist(imagePath1, imagePath2, method);
        } finally {
            IoUtil.delete(new File(imagePath1));
            IoUtil.delete(new File(imagePath2));
        }
    }

    /**
     * 获取两张图片相似度
     * @param imgPath1 第一张图片的绝对路径
     * @param imgPath2 第二张图片的绝对路径
     * @param method 比较方法
     *               Imgproc.CV_COMP_CORREL, CV_COMP_CHISQR, CV_COMP_INTERSECT, CV_COMP_BHATTACHARYYA
     *               一般选择CV_COMP_CORREL
     * */
    public static double compareHist(String imgPath1, String imgPath2, int method) {
        checkInitialized();
        Mat src1 = Imgcodecs.imread(imgPath1);
        Mat src2 = Imgcodecs.imread(imgPath2);
        Mat hvs1 = new Mat();
        Mat hvs2 = new Mat();
        //图片转HSV
        Imgproc.cvtColor(src1, hvs1,Imgproc.COLOR_BGR2HSV);
        Imgproc.cvtColor(src2, hvs2,Imgproc.COLOR_BGR2HSV);
        Mat hist1 = new Mat();
        Mat hist2 = new Mat();
        //直方图计算
        Imgproc.calcHist(Stream.of(hvs1).collect(Collectors.toList()),new MatOfInt(0),new Mat(),hist1,new MatOfInt(255) ,new MatOfFloat(0,256));
        Imgproc.calcHist(Stream.of(hvs2).collect(Collectors.toList()),new MatOfInt(0),new Mat(),hist2,new MatOfInt(255) ,new MatOfFloat(0,256));
        //图片归一化
        Core.normalize(hist1, hist1, 1, hist1.rows() , Core.NORM_MINMAX, -1, new Mat());
        Core.normalize(hist2, hist2, 1, hist2.rows() , Core.NORM_MINMAX, -1, new Mat());
        return Imgproc.compareHist(hist1, hist2, method);
    }

    /**
     * 寻找大图片中的小图片
     * @param big 大图片
     * @param small 小图片
     * @param scale 小图片宽度于大图片宽度的占比
     * @param method 比较方法
     *               Imgproc.TM_SQDIFF, TM_SQDIFF_NORMED, TM_CCORR, TM_CCORR_NORMED, TM_CCOEFF, TM_CCOEFF_NORMED
     *               一般选择TM_CCORR_NORMED
     * */
    public static MatchResult matchTemplate(BufferedImage big, BufferedImage small, double scale, int method) {
        checkInitialized();
        int originalBigWidth = big.getWidth();
        int originalSmallWidth = big.getWidth();
        if (originalSmallWidth*1.0/originalBigWidth > scale) {
            int newSmallWidth = (int) (originalBigWidth*scale);
            if (newSmallWidth != originalSmallWidth) {
                small = AwtUtil.getThumbedBufferedImage(small, newSmallWidth, small.getHeight(), true);
            }
        } else {
            int newBigWidth = (int) (originalSmallWidth/scale);
            if (newBigWidth != originalBigWidth) {
                big = AwtUtil.getThumbedBufferedImage(big, newBigWidth, big.getHeight(), true);
            }
        }
        String bigPath = oprDir + File.separator + UUID.randomUUID().toString() + ".png";
        String smallPath = oprDir + File.separator + UUID.randomUUID().toString() + ".png";
        try (OutputStream outputStream1 = new FileOutputStream(bigPath);
             OutputStream outputStream2 = new FileOutputStream(smallPath)) {
            AwtUtil.writeBufferedImage(big, "png", outputStream1);
            AwtUtil.writeBufferedImage(small, "png", outputStream2);
        } catch (Throwable e) {
            IoUtil.delete(new File(bigPath));
            IoUtil.delete(new File(smallPath));
            LogConsole.logError(OpencvApi.class, ThrowableUtil.getFullInfo(e));
            return new MatchResult();
        }
        try {
            return matchTemplate(bigPath, smallPath, method);
        } finally {
            IoUtil.delete(new File(bigPath));
            IoUtil.delete(new File(smallPath));
        }
    }

    /**
     * 寻找大图片中的小图片
     * @param big 大图片
     * @param smallList 小图片集合
     * @param scaleList 小图片宽度于大图片宽度的占比集合
     * @param similarityFlag 相似度大于等于此值便退出
     * @param method 比较方法
     *               Imgproc.TM_SQDIFF, TM_SQDIFF_NORMED, TM_CCORR, TM_CCORR_NORMED, TM_CCOEFF, TM_CCOEFF_NORMED
     *               一般选择TM_CCORR_NORMED
     * @param matchResult 匹配结果
     * @return 匹配到的small的下标
     * */
    public static int matchTemplate(BufferedImage big, List<BufferedImage> smallList, List<Double> scaleList,
                                            double similarityFlag, int method, MatchResult matchResult) {
        checkInitialized();
        if (matchResult == null) {
            throw new ExpectedException("matchResult can not be null");
        }
        if (!CommonUtil.hasElements(smallList)
                || !CommonUtil.hasElements(scaleList)
                || smallList.size() != scaleList.size()) {
            return -1;
        }
        BufferedImage originalBig = big;
        BufferedImage latestWriteBig = null;
        String bigPath = oprDir + File.separator + UUID.randomUUID().toString() + ".png";
        List<String> smallPathList = new ArrayList<>();
        try {
            int count = 0;
            for (BufferedImage small : smallList) {
                double scale = scaleList.get(count);
                int originalBigWidth = big.getWidth();
                int originalSmallWidth = big.getWidth();
                if (originalSmallWidth * 1.0 / originalBigWidth > scale) {
                    int newSmallWidth = (int) (originalBigWidth * scale);
                    if (newSmallWidth != originalSmallWidth) {
                        small = AwtUtil.getThumbedBufferedImage(small, newSmallWidth, small.getHeight(), true);
                        big = originalBig;
                    }
                } else {
                    int newBigWidth = (int) (originalSmallWidth / scale);
                    if (newBigWidth != originalBigWidth) {
                        big = AwtUtil.getThumbedBufferedImage(big, newBigWidth, big.getHeight(), true);
                    }
                }
                if (big != latestWriteBig) {
                    try (OutputStream outputStream = new FileOutputStream(bigPath)) {
                        AwtUtil.writeBufferedImage(big, "png", outputStream);
                    } catch (Throwable e) {
                        IoUtil.delete(new File(bigPath));
                        LogConsole.logError(OpencvApi.class, ThrowableUtil.getFullInfo(e));
                        count = -1;
                        break;
                    }
                    latestWriteBig = big;
                }
                String smallPath = oprDir + File.separator + UUID.randomUUID().toString() + ".png";
                smallPathList.add(smallPath);
                try (OutputStream outputStream = new FileOutputStream(smallPath)) {
                    AwtUtil.writeBufferedImage(small, "png", outputStream);
                } catch (Throwable e) {
                    IoUtil.delete(new File(smallPath));
                    LogConsole.logError(OpencvApi.class, ThrowableUtil.getFullInfo(e));
                    continue;
                }
                MatchResult oneMatchResult = matchTemplate(bigPath, smallPath, method);
                if (oneMatchResult.getSimilarity() >= similarityFlag) {
                    matchResult.setSimilarity(oneMatchResult.getSimilarity());
                    matchResult.setBigCols(oneMatchResult.getBigCols());
                    matchResult.setBigRows(oneMatchResult.getBigRows());
                    matchResult.setSmallCols(oneMatchResult.getSmallCols());
                    matchResult.setSmallRows(oneMatchResult.getSmallRows());
                    matchResult.setX(oneMatchResult.getX());
                    matchResult.setY(oneMatchResult.getY());
                    break;
                }
                count++;
            }
            if (count >= smallList.size()) {
                return -1;
            } else {
                return count;
            }
        } finally {
            IoUtil.delete(new File(bigPath));
            for (String smallPath : smallPathList) {
                IoUtil.delete(new File(smallPath));
            }
        }
    }

    /**
     * 寻找大图片中的小图片
     * @param bigPath 大图片路径
     * @param smallPath 小图片路径
     * @param method 比较方法
     *               Imgproc.TM_SQDIFF, TM_SQDIFF_NORMED, TM_CCORR, TM_CCORR_NORMED, TM_CCOEFF, TM_CCOEFF_NORMED
     *               一般选择TM_CCORR_NORMED
     * */
    public static MatchResult matchTemplate(String bigPath, String smallPath, int method) {
        checkInitialized();
        MatchResult ret = new MatchResult();
        Mat srcBig = Imgcodecs.imread(bigPath);
        Mat srcSmall = Imgcodecs.imread(smallPath);
        ret.setBigCols(srcBig.cols());
        ret.setBigRows(srcBig.rows());
        ret.setSmallCols(srcSmall.cols());
        ret.setSmallRows(srcSmall.rows());
        int resultCols = ret.getBigCols() - ret.getSmallCols() + 1;
        int resultRows = ret.getBigRows() - ret.getSmallRows() + 1;
        Mat result = new Mat(resultRows, resultCols, CvType.CV_32FC1);
        Imgproc.matchTemplate(srcBig, srcSmall, result, method);
        Core.normalize(result, result, 0, 1, Core.NORM_MINMAX, -1, new Mat());
        Point maxLoc = Core.minMaxLoc(result).maxLoc;
        ret.setX((int) maxLoc.x);
        ret.setY((int) maxLoc.y);
        String cutImgPath = oprDir + File.separator + UUID.randomUUID().toString() + ".png";
        try {
            Imgcodecs.imwrite(cutImgPath, cutImage(srcBig, new Rect((int) maxLoc.x, (int) maxLoc.y, ret.getSmallCols(), ret.getSmallRows())));
            ret.setSimilarity(compareHist(smallPath, cutImgPath, Imgproc.CV_COMP_CORREL));
            return ret;
        } finally {
            IoUtil.delete(new File(cutImgPath));
        }
    }

    /**
     * 裁剪图片
     * @param src 裁剪来源
     * @param rect 裁剪大小
     * */
    public static Mat cutImage(Mat src, Rect rect) {
        Mat ret = new Mat();
        new Mat(src, rect).copyTo(ret);
        return ret;
    }


}
